package org.jsoncrypto;

import java.math.BigInteger;

import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.agreement.ECDHBasicAgreement;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.util.BigIntegers;
import org.jsoncrypto.util.Base64;

public class JweECDH_ES extends JcLibJwe_EC {
  public JweECDH_ES(JcLibEnc enc) throws JcBaseException {
    super(enc, JcBase.KW_ALG_ECDH_ES);
  }

  public JweECDH_ES(String encrypted) throws JcBaseException {
    super(encrypted);
  }

  @Override
  public String encrypt(byte[] contentBytes, ECPrivateKeyParameters ecPrivateKeyParameters,
      ECPublicKeyParameters ecPublicKeyParameters) throws JcBaseException {
    try {
      ECDHBasicAgreement ecdhBasicAgreement = new ECDHBasicAgreement();
      ecdhBasicAgreement.init(ecPrivateKeyParameters);
      BigInteger z = ecdhBasicAgreement.calculateAgreement(ecPublicKeyParameters);
      // System.out.println("ECDH-ES z=" + z.toString());
      byte[] zBytes = BigIntegers.asUnsignedByteArray(z);
      int keylength = JcBase.enc2cekLength(mEnc.getEncStr());
      Digest kdfDigest = JcBase.getKdfDigest(mEnc.mHeader);

      byte[] secretKeyBytes = JcBase.generateCEK(zBytes, keylength / 8, kdfDigest, mEnc.getEncStr());
      SecretKeySpec secretKey = new SecretKeySpec(secretKeyBytes, "RAW");
      String encodedJweEncryptedKey = ""; 
      String encodedJweIvSegment = null;
      return mEnc.encrypt(contentBytes, secretKey, encodedJweEncryptedKey, encodedJweIvSegment );

    } catch (Exception e) {
      throw new JcBaseException(e);
    }
  }

  @Override
  public byte[] decrypt(ECPrivateKeyParameters ecPrivateKeyParameters)
      throws JcBaseException {
    try {
      ECPublicKeyParameters ecPublicKeyParameters = JcBase.getECPublicKeyParameters(mEnc.mHeader);
      
      ECDHBasicAgreement ecdhBasicAgreement = new ECDHBasicAgreement();
      ecdhBasicAgreement.init(ecPrivateKeyParameters);
      BigInteger z = ecdhBasicAgreement.calculateAgreement(ecPublicKeyParameters);
      // System.out.println("ECDH-ES z=" + z.toString());
      byte[] zBytes = BigIntegers.asUnsignedByteArray(z);
      int keylength = JcBase.enc2cekLength(mEnc.getEncStr());
      Digest kdfDigest = JcBase.getKdfDigest(mEnc.mHeader);
      byte[] secretKeyBytes = JcBase.generateCEK(zBytes, keylength / 8, kdfDigest, mEnc.getEncStr());
      mEnc.mEncodedJwtKeySegment = Base64.encodeBytes(secretKeyBytes, org.jsoncrypto.util.Base64.DONT_BREAK_LINES
          | org.jsoncrypto.util.Base64.URL);
      byte[] jwtIntegritySegmentBytes = Base64.decodeUrl(mEnc.mEncodedJwtIntegritySegment);
      
      return JcBase.doDecrypt(mEnc.mEncodedJwtHeaderSegment, mEnc.mEncodedJwtKeySegment, mEnc.mEncodedJwtCryptoSegment,
          jwtIntegritySegmentBytes, mEnc.mHeader, secretKeyBytes, mEnc.mEncodedJwtIvSegment);

    } catch (Exception e) {
      throw new JcBaseException(e);
    }
  }
}
